
*----------------------------------------------------------------
* Calculate water balance with zero initial soil moisture
*----------------------------------------------------------------
set more off

* range of years in analysis 
local yearMin = 1981
local initial_date = "01/01/`yearMin'"
local yearMax = 2012

* Parameters for water-balance model
local T_rain = 3.3
local T_snow = -10
local drofrac = 0.05
local meltmax = 0.5

* Define soil moisture and snow storage in initial period
use "..\dataRaw\county_soils_usa", clear
keep stcofips rootznaws
gen date=date("`initial_date'", "MDY") - 1
********************************
gen soil_moisture_usgs=0*rootznaws
gen soil_moisture=0*rootznaws
gen soil_moisture_corn=0*rootznaws
gen snostor=0
********************************
format date %td
save "..\temp\initial_soil_moisture", replace

* Import file of crop coefficients
import excel using "..\dataRaw\crop coefficients ET.xlsx", sheet("data") first clear
save "..\dataRaw\crop coefficients ET", replace


forvalues t = `yearMin'/`yearMax' {
		display "`t'"
qui{		
        ******************************************************************************
        * load data
        use "../dataRaw/PRISM_county_`t'", clear 
		
		gen date_str=string(month) + "/" + string(day) + "/" + string(year)
		gen date=date(date_str,"MDY")
		drop date_str
		format date %td
		xtset stcofips date

		merge m:1 stcofips using "..\dataRaw\county_soils_usa", keep(match master) keepusing(rootznaws) nogen
		
		* Snow Accumulation
		gen ppt_snow=ppt if tavg<=`T_snow'
		replace ppt_snow=ppt*(`T_rain' - tavg)/(`T_rain'-`T_snow') if tavg>`T_snow' & tavg<`T_rain'
		replace ppt_snow=0 if tavg>=`T_rain'
		gen ppt_rain=ppt-ppt_snow

		* Direct Runoff
		gen ppt_remain=ppt_rain - ppt_rain*`drofrac'

		* Snow Melt
		gen smf=0 if tavg<=`T_snow'
		replace smf=(tavg - `T_snow')/(`T_rain'-`T_snow')*`meltmax' if tavg>`T_snow' & tavg<`T_rain'
		replace smf=`meltmax' if tavg>=`T_rain'
		
		* Generate corn-specific ET
		gen plant_date=mdy(5,1,year)
		gen Days_after_plant=date-plant_date
		merge m:1 Days_after_plant using "..\dataRaw\crop coefficients ET", nogen keep(master match)
		replace kc=0.25 if kc==.
		replace kc=0.25 if kc<0.25
		gen ETc=ET0*kc
		drop plant_date Days_after_plant
		sort stcofips date
	
		gen aet=.
		gen soil_moisture=.
		gen water_deficit=.
		gen water_surplus=.
		
		gen snostor=.
		gen sm=.
		gen ppt_total=.
		gen stw=.
		gen aet_usgs=.
		gen soil_moisture_usgs=.
		gen water_deficit_usgs=.
		gen water_surplus_usgs=.	
		gen aet_corn=.
		gen soil_moisture_corn=.
		gen water_deficit_corn=.
		gen water_surplus_corn=.
		
		* append initial soil moisture and snow storage
		if `t'==`yearMin' {
			append using "..\temp\initial_soil_moisture"
		}
		else {
			local tMinus1=`t'-1
			append using `WaterBalanceDec31_`tMinus1''
		}
		
		sort stcofips date
		replace snostor=(ppt_snow + l.snostor)*(1-smf) if date>=date("01/01/`t'", "MDY")
		replace sm=(ppt_snow + l.snostor)*smf if date>=date("01/01/`t'", "MDY")
		replace ppt_total=ppt_remain + sm if date>=date("01/01/`t'", "MDY")
		
		local start = date("01/01/`t'", "MDY")
		local end = date("12/31/`t'", "MDY")
		forvalues day = `start'/`end'  {
			***********************
			* Simple Water Balance
			***********************
			replace soil_moisture=l.soil_moisture + ppt - ET0 if date==`day'
			replace soil_moisture=0 if soil_moisture<0 & date==`day'
			replace water_surplus=soil_moisture - rootznaws if soil_moisture>rootznaws & date==`day'
			replace soil_moisture=rootznaws if soil_moisture>rootznaws & date==`day'
			* If precip < ET0
			replace aet=ppt - d.soil_moisture if ppt<ET0 & date==`day'
			* If precip >=ET0
			replace aet=ET0 if ppt>=ET0 & date==`day'
			replace water_deficit=ET0-aet if aet<ET0 & date==`day'
			
			***********************
			* USGS Water Balance
			***********************
			* If precip < ET0
			replace stw=(ET0-ppt_total)*l.soil_moisture_usgs/rootznaws if ppt_total<ET0 & date==`day'
			replace aet_usgs=ppt_total + stw if ppt_total<ET0 & date==`day'
			replace soil_moisture_usgs=l.soil_moisture_usgs - stw if ppt_total<ET0 & date==`day'
			replace soil_moisture_usgs=0 if soil_moisture_usgs<0 & date==`day'
			replace water_deficit_usgs=ET0-aet_usgs if aet_usgs<ET0 & date==`day'
	
			* If precip >= ET0
			replace aet_usgs=ET0 if ppt_total>=ET0 & date==`day'
			replace soil_moisture_usgs=l.soil_moisture_usgs + ppt_total-ET0 if ppt_total>=ET0 & date==`day'
			replace water_surplus_usgs=soil_moisture_usgs - rootznaws if soil_moisture_usgs>rootznaws & date==`day'
			replace soil_moisture_usgs=rootznaws if soil_moisture_usgs>rootznaws & date==`day'
			
			***********************
			* Corn Water Balance
			***********************
			replace soil_moisture_corn=l.soil_moisture_corn + ppt - ETc if date==`day'
			replace soil_moisture_corn=0 if soil_moisture_corn<0 & date==`day'
			replace water_surplus_corn=soil_moisture_corn - rootznaws if soil_moisture_corn>rootznaws & date==`day'
			replace soil_moisture_corn=rootznaws if soil_moisture_corn>rootznaws & date==`day'
			* If precip < ETc
			replace aet_corn=ppt - d.soil_moisture_corn if ppt<ETc & date==`day'
			* If precip >=ETc
			replace aet_corn=ETc if ppt>=ETc & date==`day'
			replace water_deficit_corn=ETc-aet_corn if aet_corn<ETc & date==`day'
			
		} //end of looping over days
		replace water_deficit=0 if water_deficit==. & rootznaws!=. & ET0!=.
		replace water_surplus=0 if water_surplus==. & rootznaws!=. & ET0!=.
		replace water_deficit_usgs=0 if water_deficit_usgs==. & rootznaws!=. & ET0!=.
		replace water_surplus_usgs=0 if water_surplus_usgs==. & rootznaws!=. & ET0!=.
		replace water_deficit_corn=0 if water_deficit_corn==. & rootznaws!=. & ETc!=.
		replace water_surplus_corn=0 if water_surplus_corn==. & rootznaws!=. & ETc!=.
		
		save "..\temp\PRISM_WaterBalance_`t'_0moist", replace
		
		*keep 12/31 for initial period in next year
		keep if date==date("12/31/`t'", "MDY")
		tempfile WaterBalanceDec31_`t'
		save `WaterBalanceDec31_`t''
				
} //end of qui
} //end of looping over years





*-------------------------------------------------------------------------
* Collapse PRISM data into annual and average data for the Growing Season
*------------------------------------------------------------------------
set more off
* range of years in analysis 
local yearMin = 1981
local initial_date = "01/01/`yearMin'"
local yearMax = 2012

* define growing season season
* first month and day of month used in a year
local monthBegin = 4 
local dayMonthBegin =  1
* last month and day of month used in a year
local monthEnd   = 9 
local dayMonthEnd   = 31



* Erase the old version of the data before starting loop to append each year of data
capture erase "..\temp\PRISM_county_annual_GrowSeason_0moist.dta"

* Average data within the growing season
forvalues t = `yearMin'/`yearMax' {
		display "`t'"
qui{		
        ******************************************************************************
        * load data
        use "..\temp\PRISM_WaterBalance_`t'_0moist", clear 
		local tMinus1=`t'-1
		drop if date==date("12/31/`tMinus1'", "MDY")
		
		* keep only days for selected part of year
        qui drop if (month < `monthBegin')
        qui drop if (month > `monthEnd')
        qui drop if (month == `monthBegin') & (day < `dayMonthBegin')
        qui drop if (month == `monthEnd')   & (day > `dayMonthEnd')
		
		drop ppt_snow ppt_rain ppt_remain smf snostor sm stw
		
		* Generate Soil Moisture Deficit
		gen soil_moist_deficit_usgs=rootznaws-soil_moisture_usgs
		gen soil_moist_deficit=rootznaws-soil_moisture
		gen soil_moist_deficit_corn=rootznaws-soil_moisture_corn

        
		
		* Calculate net precipitation
		gen netppt0=ppt-ET0 
		gen netpptc=ppt-ETc 
		
		
		*****************************************************
		* Calculate the time exposed to each temperature bin
		*****************************************************
		forvalues b=-5/50{
			if (`b' < 0) {
				local b2 = abs(`b') 
				local bLabel Minus`b2' 
			} 
			else { 
				local bLabel `b' 
			}
		
			* Solve for time (in minutes) at temp b and b+1 
			* Wikipedia page for "sine wave" give the relevant formula
			* simply invert the formula to solve for time
			local bPlus1 `b'+1
			gen minb_early=asin((`b'-tmin)/(tmax-tmin))/(2*_pi*1/(2*1440))
			gen minbPlus1_early=asin((`bPlus1'-tmin)/(tmax-tmin))/(2*_pi*1/(2*1440))

			* calculate time spent greater than b degrees C and b+1 degrees C
			* divide by 1440 to convert to days
			gen tempBin`bLabel'C=2*(minbPlus1_early - minb_early)/1440
			drop minbPlus1_early minb_early
		}
		
		*****************************************************
		* Calculate the days exposed to each water deficit bin
		*****************************************************
		local interval=0.25
		forvalues j=0(`interval')10 {
			local j2=floor(`j')
			local j3=(`j'-`j2')*100
			local jLabel `j2'dot`j3'

			gen deficitBin`jLabel'=(water_deficit>=`j' & water_deficit<`j'+`interval')
			replace deficitBin`jLabel'=. if water_deficit==.
		}

		
        ******************************************************************************;
        * collapse by year
        collapse (sum) ppt dday* ET0 ET0_elev ET0_Har ETc netppt* tempBin* aet* water_deficit* water_surplus* deficitBin* ///
				(mean) soil_moist* tmin tmax tavg vpd, by(stcofips year)
		* ET0_elev was missing for some counties because elevation data was missing
		* so replace 0 with missing
		replace ET0_elev=. if ET0_elev==0
	
        ******************************************************************************;
        * append years and save data
		capture append using "..\temp\PRISM_county_annual_GrowSeason_0moist"
        sort stcofips year
        save "..\temp\PRISM_county_annual_GrowSeason_0moist", replace
} // end of quiet		
}


* Average across 1983-2012 to get normal (30-year) climate data
use "..\temp\PRISM_county_annual_GrowSeason_0moist", clear
keep if year>=1983 & year<=2012
collapse (mean) ppt dday* ET0 ET0_elev ET0_Har ETc netppt* tmin tmax tavg vpd tempBin* ///
			aet* water_deficit* water_surplus* deficitBin* soil_moist*, by(stcofips)
* ET0_elev was missing for some counties because elevation data was missing
* so replace 0 with missing
replace ET0_elev=. if ET0_elev==0
		
save "..\temp\PRISM_county_climate83-12_GrowSeason_0moist", replace


*-----------------------------------------------------------------
* Get the summary stats that I need
*-----------------------------------------------------------------
use "..\temp\PRISM_county_climate83-12_GrowSeason_0moist", clear
* ET and precipitation are measured in mm
* convert ET and precipitation to inches 
local varlist ppt ET0 ET0_elev ET0_Har ETc netppt0 netpptc ///
		aet_usgs aet aet_corn water_deficit_usgs water_deficit water_deficit_corn water_surplus_usgs water_surplus water_surplus_corn ///
		soil_moisture_usgs soil_moisture soil_moisture_corn soil_moist_deficit_usgs soil_moist_deficit soil_moist_deficit_corn		
foreach var in `varlist' {
	qui replace `var'=0.0393701*`var'
}
ren * *_0moist
rename stcofips_0moist stcofips

merge 1:1 stcofips using  "..\temp\PRISM_county_climate83-12_GrowSeason", nogen
* ET and precipitation are measured in mm
* convert ET and precipitation to inches 
local varlist ppt ET0 ET0_elev ET0_Har ETc netppt0 netpptc  ///
		aet_usgs aet aet_corn water_deficit_usgs water_deficit water_deficit_corn water_surplus_usgs water_surplus water_surplus_corn ///
		soil_moisture_usgs soil_moisture soil_moisture_corn soil_moist_deficit_usgs soil_moist_deficit soil_moist_deficit_corn		
foreach var in `varlist' {
	qui replace `var'=0.0393701*`var'
}

gen WDdiff= water_deficit_0moist - water_deficit
gen WSdiff= water_surplus_0moist - water_surplus

gen WDperc= (water_deficit_0moist - water_deficit)/water_deficit*100
gen WSperc= (water_surplus_0moist - water_surplus)/water_surplus*100

summ WD*
summ WS*
summ WDperc WSperc, detail
